#include <cassert>
#include <stdexcept>
#include <chrono>
#include "cuda_runtime.h"
#include "cuda_runtime_api.h"
#include "cublas_v2.h"
#include "cusolverDn.h"
#include "device_launch_parameters.h"
#include "common.h"
#include <omp.h>
#include <atomic>

// CUDA API error checking
#define CUDA_CHECK(err)                                                                            \
    do {                                                                                           \
        cudaError_t err_ = (err);                                                                  \
        if (err_ != cudaSuccess) {                                                                 \
            printf("CUDA error %d at %s:%d\n", err_, __FILE__, __LINE__);                          \
            throw std::runtime_error("CUDA error");                                                \
        }                                                                                          \
    } while (0)

// cusolver API error checking
#define CUSOLVER_CHECK(err)                                                                        \
    do {                                                                                           \
        cusolverStatus_t err_ = (err);                                                             \
        if (err_ != CUSOLVER_STATUS_SUCCESS) {                                                     \
            printf("cusolver error %d at %s:%d\n", err_, __FILE__, __LINE__);                      \
            throw std::runtime_error("cusolver error");                                            \
        }                                                                                          \
    } while (0)

// cublas API error checking
#define CUBLAS_CHECK(err)                                                                          \
    do {                                                                                           \
        cublasStatus_t err_ = (err);                                                               \
        if (err_ != CUBLAS_STATUS_SUCCESS) {                                                       \
            printf("cublas error %d at %s:%d\n", err_, __FILE__, __LINE__);                        \
            throw std::runtime_error("cublas error");                                              \
        }                                                                                          \
    } while (0)

template<typename data_t>
void measure_latency(int len, int warm_cnt = 300, int test_cnt = 1000)
{
    cudaEvent_t start, stop;
    float tms;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));
    // data_t* hbuf = new data_t [len];
    data_t* hbuf; CUDA_CHECK(cudaHostAlloc((void**)&hbuf, len * sizeof(*hbuf), cudaHostAllocDefault));
    for (int i = 0; i < len; i++) hbuf[i] = i;

    data_t* dbuf = nullptr;
    const size_t bytes = sizeof(data_t) * len;
    printf("Bytes %zu\n", bytes);
    CUDA_CHECK(cudaMalloc((void**)&dbuf, bytes));
    
    // 测试H->D
    for (int cnt = 0; cnt < warm_cnt; cnt++) {
        cudaMemcpy(dbuf, hbuf, bytes, cudaMemcpyHostToDevice);
        // hbuf += 1.0;
    }
    CUDA_CHECK(cudaEventRecord(start));
    auto beg = std::chrono::high_resolution_clock::now();

    for (int cnt = 0; cnt < test_cnt; cnt++) {
        cudaMemcpy(dbuf, hbuf, bytes, cudaMemcpyHostToDevice);
        // hbuf += 1.0;
    }
    auto end = std::chrono::high_resolution_clock::now();
    auto E2E = (std::chrono::duration_cast<std::chrono::nanoseconds>(end - beg)).count() * 1e-9;
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));
    CUDA_CHECK(cudaEventElapsedTime(&tms, start, stop));
    printf("H->D: Total %.3f ms of %d cnt. Single %.3f ms. Val %.2f\n", tms, test_cnt, tms / test_cnt, hbuf);
    printf("E2E %.6f s\n", E2E);
    // 测试D->H
    for (int cnt = 0; cnt < warm_cnt; cnt++) {
        cudaMemcpy(hbuf, dbuf, bytes, cudaMemcpyDeviceToHost);
        // hbuf += 1.0;
    }
    CUDA_CHECK(cudaEventRecord(start));
    beg = std::chrono::high_resolution_clock::now();

    for (int cnt = 0; cnt < test_cnt; cnt++) {
        cudaMemcpy(hbuf, dbuf, bytes, cudaMemcpyDeviceToHost);
        // hbuf += 1.0;
    }

    end = std::chrono::high_resolution_clock::now();
    E2E = (std::chrono::duration_cast<std::chrono::nanoseconds>(end - beg)).count() * 1e-9;
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));
    CUDA_CHECK(cudaEventElapsedTime(&tms, start, stop));
    printf("D->H: Total %.3f ms of %d cnt. Single %.3f ms. Val %.2f\n", tms, test_cnt, tms / test_cnt, hbuf);
    printf("E2E %.6f s\n", E2E);

    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));
    CUDA_CHECK(cudaFree(dbuf));
    // delete[] hbuf;
    CUDA_CHECK(cudaFreeHost(hbuf));
}


int main(int argc, char** argv)
{
    measure_latency<double>(10000000, 10, 20);

    measure_latency<float>(10000000, 10, 20);

    return 0;
}
